#include "jerasure.h"

static int create_r6_min_density_matrix( int k, int w, JER_Matrix *matrix );
static int is_prime( int p );

static int m4[25] = { 4, 10, 12, 16, 18, 22, 28, 36, 40, 42, 46, 52, 58, 
					  60, 66, 70, 78, 82, 88, 96, 100, 102, 106, 108, 112 };
static int m5[22] = { 4, 10, 12, 18, 22, 28, 36, 40, 46, 52, 58, 60, 66, 
					  70, 78, 82, 96, 100, 102, 106, 108, 112 };
static int m6[20] = { 10, 18, 22, 28, 36, 40, 46, 52, 58, 60, 66, 70, 78, 
					  82, 96, 100, 102, 106, 108, 112 };
static int m7[15] = { 18, 28, 36, 46, 52, 58, 60, 66, 70, 78, 82, 96, 100,
					  102, 106 };
static int m8[9] = { 36, 46, 52, 58, 60, 66, 82, 100, 106 };
static int eotablesize[5] = { 25, 22, 20, 15, 9 };
static int *eotable[5] = { m4, m5, m6, m7, m8 };

int R6_Min_Density_Smallest_W(int k){
	int w;
	for(w=k; !is_prime(w) && !is_prime( w + 1 ) && w != 8;w++);
	return w;
}

JER_Gen_T *R6_Min_Density_Generator(int k, int w){

	JER_Gen_T *gen;
	JER_Matrix *matrix;

	gen = new JER_Gen_T();
	gen->N = k + 2;
	gen->K = k;
	gen->WPD = w;

	gen->Systematic = true;
	gen->PDrive = true;

	matrix = new JER_Matrix( 2*w, k*w, 1 );

	if( create_r6_min_density_matrix( k, w, matrix ) )
	{
		delete matrix;
		delete gen;
		return NULL;
	}
	else
	{
		gen->M = matrix;
		return gen;
	}
}

int R6_Min_Density_Generator(int k, int w, JER_Gen_T &g){

	JER_Matrix *matrix;
	g.N = k + 2;
	g.K = k;
	g.WPD = w;

	g.Systematic = true;
	g.PDrive = true;

	//get rid of old schedules
	g.Delete_Schedules();

	matrix = new JER_Matrix( 2*w, k*w, 1 );

	if( create_r6_min_density_matrix( k, w, matrix ) )
	{
		delete matrix;
		g.M = NULL;
		return -1;
	}
	else
	{
		g.M = matrix;
		return 0;
	}

}

int create_r6_min_density_matrix( int k, int w, JER_Matrix *matrix )
{
	int i, j, l, r, c, m, p;
	int super_c;
	/* Create identity row. */
	for( i = 0; i < w; i++ )
	{
		for( j = i; j < k * w; j += w )
			matrix->Set( i, j, 1 );
	}

	/* Liberation codes have w prime. */
	if( is_prime( w ) )
	{
		for( super_c = 0; super_c < k; super_c++ )
		{
			for( j = 0; j < w; j++ )
			{
				l = (super_c+j)%w+super_c*w;
				matrix->Set( j + w, l, 1 );
			}
			if( super_c > 0 )
			{
				l = (super_c*((w-1)/2))%w;
				matrix->Set( w+l, super_c*w+(l+super_c-1)%w, 1 );
			}
		}
	}
	/* Blaum-Roth codes have w+1 prime. */
	else if( is_prime( w + 1 ) )
	{
		p = w + 1;

		for( super_c = 0; super_c < k; super_c++ )
		{
			c = super_c * w;
			r = w;

			if( super_c == 0 )
			{
				for( c = 0; c < w; c++ )
				{
					matrix->Set( r, c, 1 );
					r++;
				}
			}
			else
			{
				i = super_c;
				for (l = 1; l <= w; l++) {
					if (l != p-i) {
						m = l+i;
						if (m >= p) m -= p;
						m--;
						matrix->Set(r,c+m,1);
					} else {
						matrix->Set(r,c+i-1,1);
						if (i%2 == 0) {
							m = i/2;
						} else {
							m = (p/2) + 1 + (i/2);
						}
						m--;
						matrix->Set(r,c+m,1);
					}
					r++;
				}
			}
		}
	}
	/* Liber8tion codes have w = 8. */
	else if( w == 8 )
	{
		super_c = 0;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+0,1);
			matrix->Set(w+1,super_c*w+1,1);
			matrix->Set(w+2,super_c*w+2,1);
			matrix->Set(w+3,super_c*w+3,1);
			matrix->Set(w+4,super_c*w+4,1);
			matrix->Set(w+5,super_c*w+5,1);
			matrix->Set(w+6,super_c*w+6,1);
			matrix->Set(w+7,super_c*w+7,1);
		}

		super_c=1;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+7,1);
			matrix->Set(w+1,super_c*w+3,1);
			matrix->Set(w+2,super_c*w+0,1);
			matrix->Set(w+3,super_c*w+2,1);
			matrix->Set(w+4,super_c*w+6,1);
			matrix->Set(w+4,super_c*w+7,1);
			matrix->Set(w+5,super_c*w+1,1);
			matrix->Set(w+6,super_c*w+5,1);
			matrix->Set(w+7,super_c*w+4,1);
		}

		super_c=2;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+6,1);
			matrix->Set(w+1,super_c*w+2,1);
			matrix->Set(w+1,super_c*w+3,1);
			matrix->Set(w+2,super_c*w+4,1);
			matrix->Set(w+3,super_c*w+0,1);
			matrix->Set(w+4,super_c*w+7,1);
			matrix->Set(w+5,super_c*w+3,1);
			matrix->Set(w+6,super_c*w+1,1);
			matrix->Set(w+7,super_c*w+5,1);
		}

		super_c=3;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+2,1);
			matrix->Set(w+1,super_c*w+5,1);
			matrix->Set(w+2,super_c*w+7,1);
			matrix->Set(w+3,super_c*w+6,1);
			matrix->Set(w+4,super_c*w+0,1);
			matrix->Set(w+5,super_c*w+3,1);
			matrix->Set(w+5,super_c*w+4,1);
			matrix->Set(w+6,super_c*w+4,1);
			matrix->Set(w+7,super_c*w+1,1);
		}

		super_c=4;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+5,1);
			matrix->Set(w+1,super_c*w+6,1);
			matrix->Set(w+2,super_c*w+0,1);
			matrix->Set(w+2,super_c*w+1,1);
			matrix->Set(w+3,super_c*w+7,1);
			matrix->Set(w+4,super_c*w+2,1);
			matrix->Set(w+5,super_c*w+4,1);
			matrix->Set(w+6,super_c*w+3,1);
			matrix->Set(w+7,super_c*w+0,1);
		}

		super_c=5;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+1,1);
			matrix->Set(w+1,super_c*w+2,1);
			matrix->Set(w+2,super_c*w+3,1);
			matrix->Set(w+3,super_c*w+4,1);
			matrix->Set(w+4,super_c*w+5,1);
			matrix->Set(w+5,super_c*w+6,1);
			matrix->Set(w+6,super_c*w+7,1);
			matrix->Set(w+7,super_c*w+0,1);
			matrix->Set(w+7,super_c*w+2,1);
		}

		super_c=6;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+3,1);
			matrix->Set(w+1,super_c*w+0,1);
			matrix->Set(w+2,super_c*w+6,1);
			matrix->Set(w+3,super_c*w+5,1);
			matrix->Set(w+4,super_c*w+1,1);
			matrix->Set(w+5,super_c*w+7,1);
			matrix->Set(w+6,super_c*w+4,1);
			matrix->Set(w+6,super_c*w+5,1);
			matrix->Set(w+7,super_c*w+2,1);
		}

		super_c=7;
		if (k > super_c)
		{
			matrix->Set(w+0,super_c*w+4,1);
			matrix->Set(w+1,super_c*w+7,1);
			matrix->Set(w+2,super_c*w+1,1);
			matrix->Set(w+3,super_c*w+5,1);
			matrix->Set(w+3,super_c*w+1,1);
			matrix->Set(w+4,super_c*w+3,1);
			matrix->Set(w+5,super_c*w+2,1);
			matrix->Set(w+6,super_c*w+0,1);
			matrix->Set(w+7,super_c*w+6,1);
		}
	}
	else
		return -1;
	return 0;
}

static int is_prime( int p )
{
	int i;
	if( p < 40 )
	{
		if( p == 2 || p == 3 || p == 5 || p == 7 ||
			p == 11 || p == 13 || p == 17 || p == 19 ||
			p == 23 || p == 29 || p == 31 || p == 37 )
			return 1;
		else return 0;
	}
	else
	{
		for( i = 2; i*i < p; i++ )
		{
			if( p % i == 0 )
				return 0;
		}
		return 1;
	}
}

int Gen_Evenodd_Smallest_W(int k, int m){

	int w;
	int tbl_size;
	int *tbl;
	int i;

	if(m == 2){
		//regular raid 6 evenodd coding	
		//return smallest w >= k-1, w+1 is prime
		for(w=k-1;!is_prime(w+1);w++);
		return w;
	}else if(m >= 4 && m<=8){
		//generalized, use table
		tbl_size = eotablesize[m-4];
		tbl = eotable[m-4];
		for(i=0;i<tbl_size;i++){
			if(tbl[i] >= k-1){
				return tbl[i];
			}
		}
		return -1;
	}else{
		return -1;
	}
}

void fill_geo_matrix(int k, int m, int w, JER_Matrix *jm){

	int prim, i, gen, j, exp;
	int r,c;

	for (r = 0; r < m; r++) {
		for (c = 0; c < k; c++) {
			exp = r*c;

			prim = 0;
			for (i = 0; i <= w; i++) prim |= (1 << i);
			exp %= (w+1);

			gen = (exp == w) ? (prim ^ (1 << w)) : (1 << exp);
			for (i = 0; i < w; i++) {
				for (j = 0; j < w; j++) {
					if (gen & (1 << j)) jm->Set(r*w+j, c*w+i, 1);
				}
				gen <<= 1;
				if (gen & (1 << w)) gen ^= prim;
			}
		}
	}
}
//used to create gen_evenodd generator
JER_Gen_T *Gen_Evenodd_Generator(int k, int m, int w){

	JER_Gen_T *g;
	JER_Matrix *jm;

	g= new JER_Gen_T;
	g->N = m+k;
	g->K = k;
	g->WPD = w;
	g->Systematic = true;
	g->PDrive = true;

	jm = new JER_Matrix(m*w,k*w,1);
	g->M = jm;

	fill_geo_matrix(k,m,w,jm);

	return g;

}
int Gen_Evenodd_Generator(int k, int m, int w, JER_Gen_T &g){

	int prim, i, gen, j, exp;
	int r,c;
	JER_Matrix *jm;

	//error check
	if(is_prime(w+1) == 0 || k > (w+1)){
		return -1;	
	}

	jm = g.M;

	if(jm == NULL){
		jm = new JER_Matrix(m*w,k*w,1);
		g.M = jm;
	}else{
		jm->Create_Empty(m*w,k*w,1);
	}

	g.N = m+k;
	g.K = k;
	g.WPD = w;
	g.Systematic = true;
	g.PDrive = true;
	g.rs_r6 = false;

	fill_geo_matrix(k,m,w,jm);

	//get rid of old schedules
	g.Delete_Schedules();

	return 0;

}

void fill_rdp_matrix(int k, int m, int w, JER_Matrix *jm){

	int i, j, x, y, z, l, p, r;

	p = w+1;//should be prime

	/* First put the identity matrices for drive 0 */
	for (i = 0; i < w; i++) {
		for (x = 0; x < k; x++) {
			jm->Set(i, x*w+i, 1);
		}
	}

	/* Now do the rest of the drives. */
	for (j = 1; j < m; j++) {
		for (i = 0; i < w; i++) {  /* From p. 2785, we're defining a_i,{p-1+j} */
			/* Ok: a_{i,p-1+j} = sum_{l=0}^{p-1} a_{i-j*l,l}, The first subscript
			   is the row taken modulo p, with the proviso that row p-1 consists
			   of all zeros.  This means that a_{p-1,p-1} = 0, since the parities
			   of all the phantom elements is zero */
			for (l = 0; l < (w+1); l++) {
				r = i-j*l;
				if (r >= 0) {
					r %= p;
				} else {
					r = (-r) % p;
					if (r != 0) r = p-r;
				}
				if (r < p-1 && l < w) {
					if (l < k) {
						jm->Set(j*w+i, l*w+r, 1);
					}
				} else if (r < p-1 && l == w) {
					for (x = 0; x < k*w; x++) jm->Set(j*w+i, x, jm->Get(r, x) ^ jm->Get(j*w+i, x));
				}
			}
		}
	}

}

JER_Gen_T *Gen_RDP_Generator(int k, int m, int w){

	JER_Gen_T *g;
	JER_Matrix *jm;

	//error check
	if(is_prime(w+1) == 0 || k >w){
		return NULL;	
	}

	g= new JER_Gen_T;
	g->N = m+k;
	g->K = k;
	g->WPD = w;
	g->Systematic = true;
	g->PDrive = true;

	jm = new JER_Matrix(m*w,k*w,1);
	g->M = jm;

	fill_rdp_matrix(k,m,w,jm);

	return g;

}
int Gen_RDP_Generator(int k, int m, int w, JER_Gen_T &g){

	JER_Matrix *jm;

	//error check
	if(is_prime(w+1) == 0 || k >w){
		return -1;	
	}

	jm = g.M;

	if(jm == NULL){
		jm = new JER_Matrix(m*w,k*w,1);
		g.M = jm;
	}else{
		jm->Create_Empty(m*w,k*w,1);
	}

	g.N = m+k;
	g.K = k;
	g.WPD = w;
	g.Systematic = true;
	g.PDrive = true;
	g.rs_r6 = false;

	fill_rdp_matrix(k,m,w,jm);

	//get rid of old schedules
	g.Delete_Schedules();

	return 0;

}
